import os
import pandas as pd
import matplotlib.pyplot as plt

# 设置CSV文件路径和图片保存路径
csv_file_path = 'rsei_output.csv'  # 替换为你的CSV文件路径
pic_path = 'path_to_save_image'  # 替换为保存图片的路径

# 确保图片保存路径存在
os.makedirs(pic_path, exist_ok=True)

# 读取CSV文件
dataset = pd.read_csv(csv_file_path)

# 确保data是dataset的一个副本
data = dataset[['NDBSI', 'NDVI', 'WET', 'LST', 'RSEI']].copy()
data.loc[:, 'grade'] = pd.cut(x=data['RSEI'], bins=[0, 0.2, 0.4, 0.6, 0.8, 1.0], labels=[1, 2, 3, 4, 5])

# 直接从dataset中提取'latitude'和'longitude'列，并添加到data中
data = pd.concat([data, dataset[['latitude', 'longitude']]], axis=1)

# 自定义颜色映射，从红色到绿色
cmap = plt.get_cmap('RdYlGn')  # 使用RdYlGn颜色映射，红色到绿色
norm = plt.Normalize(vmin=0, vmax=1)  # 归一化RSEI值到0-1范围

# 绘制热力图
plt.figure(figsize=(10, 8))
plt.scatter(data['latitude'], data['longitude'], c=data['RSEI'], cmap=cmap, norm=norm, s=10)
plt.colorbar(label='RSEI')  # 添加颜色条
plt.title('RSEI Heatmap')
plt.xlabel('Latitude')
plt.ylabel('Longitude')

# 构建图片的完整路径
file_name = os.path.splitext(os.path.basename(csv_file_path))[0]
image_path = os.path.join(pic_path, f'{file_name}.png')

# 保存图片
plt.savefig(image_path, dpi=1200)

# 显示图片
plt.show()